'''
https://leetcode.cn/problems/distribute-coins-in-binary-tree/description/
'''
from typing import Optional


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:

        def f(root):
            if not root:
                # 当前子树节点的 1数量 2多少枚硬币 3走了多少步
                return 0, 0, 0
            lnodes, lcoins, lsteps = f(root.left)
            rnodes, rcoins, rsteps = f(root.right)
            nodes = lnodes + rnodes + 1
            coins = lcoins + rcoins + root.val
            steps = rsteps + lsteps + abs(lnodes - lcoins) + abs(rnodes - rcoins)
            return nodes, coins, steps
        return f(root)[2]
